"""
GPU配置优化文件 - 针对BGE-M3模型的GPU加速设置

注意：此文件已弃用，因为BGE-M3现在通过API调用，无需本地GPU。
保留此文件以保持向后兼容性，但实际功能已迁移到API调用模式。
"""

import torch
import logging

logger = logging.getLogger(__name__)

def check_gpu_availability():
    """
    检查GPU可用性和CUDA支持
    
    Returns:
        dict: GPU信息字典
    """
    gpu_info = {
        'cuda_available': torch.cuda.is_available(),
        'device_count': 0,
        'current_device': None,
        'device_name': None,
        'memory_total': 0,
        'memory_free': 0
    }
    
    if torch.cuda.is_available():
        gpu_info['device_count'] = torch.cuda.device_count()
        gpu_info['current_device'] = torch.cuda.current_device()
        gpu_info['device_name'] = torch.cuda.get_device_name()
        
        # 获取内存信息（MB）
        memory_total = torch.cuda.get_device_properties(0).total_memory / 1024**2
        memory_free = (torch.cuda.get_device_properties(0).total_memory - 
                      torch.cuda.memory_allocated()) / 1024**2
        
        gpu_info['memory_total'] = int(memory_total)
        gpu_info['memory_free'] = int(memory_free)
    
    return gpu_info

def get_optimal_device():
    """
    获取最优设备配置
    
    Returns:
        str: 推荐的设备类型 ('cuda', 'mps', 'cpu')
    """
    # 检查CUDA
    if torch.cuda.is_available():
        gpu_info = check_gpu_availability()
        # 至少需要2GB GPU内存来运行BGE-M3
        if gpu_info['memory_free'] > 2048:  
            logger.info(f"检测到可用GPU: {gpu_info['device_name']}")
            logger.info(f"GPU内存: {gpu_info['memory_free']}MB 可用")
            return 'cuda'
        else:
            logger.warning(f"GPU内存不足({gpu_info['memory_free']}MB)，建议使用CPU")
    
    # 检查Apple Silicon MPS
    if hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
        logger.info("检测到Apple Silicon GPU，使用MPS加速")
        return 'mps'
    
    logger.info("使用CPU进行计算")
    return 'cpu'

def create_optimized_embeddings_model():
    """
    创建BGE-M3嵌入模型（通过API调用）
    
    注意：此函数已更新为使用API调用，不再需要本地模型或GPU。
    
    Returns:
        BGEM3APIEmbeddings: API调用的嵌入模型
    """
    from bge_m3_api_embeddings import BGEM3APIEmbeddings
    
    logger.info("使用BGE-M3 API嵌入模型（无需本地模型或GPU）")
    
    try:
        embeddings_model = BGEM3APIEmbeddings()
        logger.info(f"✅ 成功创建BGE-M3嵌入模型 (API调用模式)")
        logger.info("API调用模式：无需本地模型，无需GPU，自动处理")
        return embeddings_model
        
    except Exception as e:
        logger.error(f"❌ 创建嵌入模型失败: {str(e)}")
        logger.error("请检查:")
        logger.error("1. 是否已设置环境变量 BGE_M3_API_KEY 或 MAAS_API_KEY")
        logger.error("2. 是否已设置环境变量 BGE_M3_API_URL（可选，有默认值）")
        logger.error("3. 网络连接是否正常")
        raise

def optimize_for_large_documents():
    """
    针对大型文档处理的优化建议
    
    Returns:
        dict: 优化配置建议
    """
    device = get_optimal_device()
    
    config = {
        'chunk_size': 1000,
        'chunk_overlap': 200,
        'batch_processing': True,
        'max_batch_size': 8
    }
    
    if device == 'cuda':
        # GPU可以处理更大的分块和批次
        config.update({
            'chunk_size': 1500,
            'chunk_overlap': 300,
            'max_batch_size': 32
        })
    elif device == 'mps':
        config.update({
            'chunk_size': 1200,
            'chunk_overlap': 250,
            'max_batch_size': 16
        })
    
    return config

if __name__ == "__main__":
    """测试GPU配置"""
    print("🔧 GPU配置检测")
    print("=" * 30)
    
    # 检查GPU信息
    gpu_info = check_gpu_availability()
    print(f"CUDA可用: {gpu_info['cuda_available']}")
    
    if gpu_info['cuda_available']:
        print(f"GPU数量: {gpu_info['device_count']}")
        print(f"GPU型号: {gpu_info['device_name']}")
        print(f"总显存: {gpu_info['memory_total']} MB")
        print(f"可用显存: {gpu_info['memory_free']} MB")
    
    # 获取推荐设备
    device = get_optimal_device()
    print(f"\n推荐设备: {device}")
    
    # 显示优化配置
    config = optimize_for_large_documents()
    print(f"\n优化配置建议:")
    for key, value in config.items():
        print(f"  {key}: {value}")
    
    print(f"\n✅ GPU配置检测完成")
    
    # 可选：测试模型创建（注释掉以避免实际下载）
    # print(f"\n🧪 测试模型创建...")
    # model = create_optimized_embeddings_model()
    # print(f"模型创建成功！")

